import type MarkdownIt from "markdown-it";
import type { KatexOptions } from "katex";
import katex from "katex";

type KatexPluginOptions = KatexOptions & {
  throwOnError?: boolean;
};

/** Checks if a `$` can be opening/closing delimiter (copied from markdown-it-katex). */
function isValidDelim(state: any, pos: number) {
  const prevChar = pos > 0 ? state.src.charCodeAt(pos - 1) : -1;
  const nextChar = pos + 1 <= state.posMax ? state.src.charCodeAt(pos + 1) : -1;

  const canClose = !(
    prevChar === 0x20 ||
    prevChar === 0x09 ||
    (nextChar >= 0x30 && nextChar <= 0x39)
  );

  const canOpen = !(nextChar === 0x20 || nextChar === 0x09);

  return { can_open: canOpen, can_close: canClose };
}

function mathInline(state: any, silent: boolean) {
  if (state.src[state.pos] !== "$") return false;

  const start = state.pos + 1;
  let match = start;

  const firstDelimCheck = isValidDelim(state, state.pos);
  if (!firstDelimCheck.can_open) {
    if (!silent) state.pending += "$";
    state.pos += 1;
    return true;
  }

  let pos = match;
  while ((match = state.src.indexOf("$", match)) !== -1) {
    pos = match - 1;
    while (state.src[pos] === "\\") pos -= 1;
    if (((match - pos) % 2) === 1) break;
    match += 1;
  }

  if (match === -1) {
    if (!silent) state.pending += "$";
    state.pos = start;
    return true;
  }

  if (match - start === 0) {
    if (!silent) state.pending += "$$";
    state.pos = start + 1;
    return true;
  }

  const closingDelimCheck = isValidDelim(state, match);
  if (!closingDelimCheck.can_close) {
    if (!silent) state.pending += "$";
    state.pos = start;
    return true;
  }

  if (!silent) {
    const token = state.push("math_inline", "math", 0);
    token.markup = "$";
    token.content = state.src.slice(start, match);
  }

  state.pos = match + 1;
  return true;
}

function mathBlock(state: any, start: number, end: number, silent: boolean) {
  let pos = state.bMarks[start] + state.tShift[start];
  const max = state.eMarks[start];

  if (pos + 2 > max || state.src.slice(pos, pos + 2) !== "$$") return false;

  pos += 2;
  let firstLine = state.src.slice(pos, max);

  if (silent) return true;

  let found = false;
  let next = start;
  let lastLine;
  let lastPos;

  if (firstLine.trim().slice(-2) === "$$") {
    firstLine = firstLine.trim().slice(0, -2);
    found = true;
  }

  while (!found) {
    next += 1;
    if (next >= end) break;

    pos = state.bMarks[next] + state.tShift[next];
    const lineMax = state.eMarks[next];

    if (pos < lineMax && state.tShift[next] < state.blkIndent) break;

    if (state.src.slice(pos, lineMax).trim().slice(-2) === "$$") {
      lastPos = state.src.slice(0, lineMax).lastIndexOf("$$");
      lastLine = state.src.slice(pos, lastPos);
      found = true;
    }
  }

  state.line = next + 1;

  const token = state.push("math_block", "math", 0);
  token.block = true;
  token.content =
    (firstLine && firstLine.trim() ? `${firstLine}\n` : "") +
    state.getLines(start + 1, next, state.tShift[start], true) +
    (lastLine && lastLine.trim() ? lastLine : "");
  token.map = [start, state.line];
  token.markup = "$$";
  return true;
}

export default function markdownItKatex(
  md: MarkdownIt,
  options: KatexPluginOptions = {}
) {
  const baseOptions = { ...options };

  const renderKatex = (latex: string, displayMode: boolean) => {
    try {
      return katex.renderToString(latex, { ...baseOptions, displayMode });
    } catch (error) {
      if (baseOptions.throwOnError) {
        console.error(error);
      }
      return latex;
    }
  };

  md.inline.ruler.after("escape", "math_inline", mathInline);
  md.block.ruler.after("blockquote", "math_block", mathBlock, {
    alt: ["paragraph", "reference", "blockquote", "list"],
  });

  md.renderer.rules.math_inline = (tokens: any, idx: number) =>
    renderKatex(tokens[idx].content, false);

  md.renderer.rules.math_block = (tokens: any, idx: number) =>
    `${renderKatex(tokens[idx].content, true)}\n`;
}

